import pandas as pd

BERT_PATH = "/home/jie/.cache/modelscope/hub/tiansz/bert-base-chinese"
BERT_train_file = "../data/bert_train_dataset.csv"

cls_name = [
    "光纤光缆、光器件/光模块、光系统设备等光通信",
    "芯片、半导体、存储器等集成电路",
    "显示器、液晶面板、触摸屏、虚拟现实、VR等显示",
    "手机、平板电脑、智能家居等智能终端",
    "5G网络",
    "互联网、云计算、大数据、物联网等",
    "传统汽车整车制造",
    "传统汽车关键核心零部件",
    "新能源整车制造",
    "新能源汽车关键核心零部件",
    "智能网联汽车配套产业",
    "汽车后市场服务",
    "医药流通",
    "生物制药、服务、制造等生物医药",
    "医疗器械",
    "生物农业",
    "健康服务",
    "3D打印与激光加工装备",
    "机器人",
    "高档数控机床",
    "绿色智能船舶、高技术船舶与海洋工程装备",
    "轨道交通装备",
    "工程机械装备",
    "高端能源装备",
    "金属新材料（含钢铁）",
    "石化化工新材料（含石化）",
    "无机非金属材料",
    "高性能纤维及制品和复合材料",
    "未来、前沿新材料",
    "电力、交通、建筑工程设计",
    "智能建造",
    "新零售、电子商务",
    "快递服务",
    "航空物流",
    "临港物流",
    "保税物流",
    "冷链物流",
    "货币金融服务",
    "证券服务",
    "投资服务",
    "保险业",
    "商务服务",
    "先进环保",
    "高效节能",
    "资源综合利用",
    "新型储能",
    "安全应急",
    "绿色消费品（含食品、纺织）",
    "滨江滨湖旅游",
    "夜色精品旅游",
    "生态休闲旅游",
    "红色文化旅游",
    "旅游演艺",
    "网络安全硬件",
    "网络安全软件（含元宇宙、区块链）",
    "网络安全服务",
    "航空维修及改装",
    "通航运营服务",
    "运载火箭及发射服务",
    "卫星制造及应用",
    "北斗基础构件",
    "北斗平台终端",
    "北斗应用服务",
    "人工智能芯片",
    "计算机视觉、自然语言处理、机器学习等AI技术",
    "无人机",
    "智能可穿戴设备",
    "智能机器人",
    "通用人工智能",
    "动漫设计",
    "游戏设计",
    "电竞直播",
    "设计服务",
    "数字传媒",
    "会议会展",
    "制氢、储氢及加氢等氢能设备",
    "电磁装备制造",
    "高速轨道交通",
    "量子通信",
    "量子计算",
    "量子测量",
    "超级计算",
    "类脑计算",
    "脑疾病诊疗",
    "脑机接口",
    "深地探测",
    "深海探测",
    "深空观测",
]

raw_attr = [
    "光通信",
    "集成电路",
    "新型显示（含未来显示）",
    "智能终端",
    "5G网络（含未来网络）",
    "互联网与云计算、大数据服务",
    "传统整车制造",
    "传统汽车关键核心零部件",
    "新能源整车制造",
    "新能源汽车关键核心零部件",
    "智能网联汽车配套产业",
    "汽车后市场服务",
    "医药流通",
    "生物医药",
    "医疗器械",
    "生物农业",
    "健康服务",
    "3D打印与激光加工装备",
    "机器人",
    "高档数控机床",
    "绿色智能船舶",
    "轨道交通装备",
    "工程机械装备",
    "高端能源装备",
    "金属新材料（含钢铁）",
    "石化化工新材料（含石化）",
    "无机非金属材料",
    "高性能纤维及制品和复合材料",
    "未来新材料",
    "工程设计",
    "智能建造",
    "电子商务",
    "快递服务",
    "航空物流",
    "临港物流",
    "保税物流",
    "冷链物流",
    "货币金融服务",
    "证券服务",
    "投资服务",
    "保险业",
    "商务服务",
    "先进环保",
    "高效节能",
    "资源综合利用",
    "新型储能",
    "安全应急",
    "绿色消费品（含食品、纺织）",
    "滨江滨湖旅游",
    "夜色精品旅游",
    "生态休闲旅游",
    "红色文化旅游",
    "旅游演艺",
    "网络安全硬件",
    "网络安全软件（含元宇宙、区块链）",
    "网络安全服务",
    "航空维修及改装",
    "通航运营服务",
    "运载火箭及发射服务",
    "卫星制造及应用",
    "北斗基础构件",
    "北斗平台终端",
    "北斗应用服务",
    "人工智能芯片",
    "软件与AI技术",
    "无人机",
    "智能可穿戴设备",
    "智能机器人",
    "通用人工智能",
    "动漫设计",
    "游戏设计",
    "电竞直播",
    "设计服务",
    "数字传媒",
    "会议会展",
    "氢能",
    "电磁装备制造",
    "高速轨道交通",
    "量子通信",
    "量子计算",
    "量子测量",
    "超级计算",
    "类脑计算",
    "脑疾病诊疗",
    "脑机接口",
    "深地探测",
    "深海探测",
    "深空观测",
]

wuhan_cls_struct = {
    "“光芯屏端网”新一代信息技术": [
        "光通信",
        "集成电路",
        "新型显示（含未来显示）",
        "智能终端",
        "5G网络（含未来网络）",
        "互联网与云计算、大数据服务",
    ],
    "汽车制造和服务": [
        "传统整车制造",
        "传统汽车关键核心零部件",
        "新能源整车制造",
        "新能源汽车关键核心零部件",
        "智能网联汽车配套产业",
        "汽车后市场服务",
    ],
    "大健康和生物技术": [
        "医药流通",
        "生物医药",
        "医疗器械",
        "生物农业",
        "健康服务",
    ],
    "高端装备和先进基础材料": [
        "3D打印与激光加工装备",
        "机器人",
        "高档数控机床",
        "绿色智能船舶",
        "轨道交通装备",
        "工程机械装备",
        "高端能源装备",
        "金属新材料（含钢铁）",
        "石化化工新材料（含石化）",
        "无机非金属材料",
        "高性能纤维及制品和复合材料",
        "未来新材料",
    ],
    "智能建造": [
        "工程设计",
        "智能建造",
    ],
    "商贸物流": [
        "电子商务",
        "快递服务",
        "航空物流",
        "临港物流",
        "保税物流",
        "冷链物流",
    ],
    "现代金融": [
        "货币金融服务",
        "证券服务",
        "投资服务",
        "保险业",
        "商务服务",
    ],
    "绿色环保": [
        "先进环保",
        "高效节能",
        "资源综合利用",
        "新型储能",
        "安全应急",
        "绿色消费品（含食品、纺织）",
    ],
    "文化旅游": [
        "滨江滨湖旅游",
        "夜色精品旅游",
        "生态休闲旅游",
        "红色文化旅游",
        "旅游演艺",
    ],
    "网络安全": [
        "网络安全硬件",
        "网络安全软件（含元宇宙、区块链）",
        "网络安全服务",
    ],
    "航空航天": [
        "航空维修及改装",
        "通航运营服务",
        "运载火箭及发射服务",
        "卫星制造及应用",
    ],
    "空天信息": [
        "北斗基础构件",
        "北斗平台终端",
        "北斗应用服务",
    ],
    "人工智能": [
        "人工智能芯片",
        "软件与AI技术",
        "无人机",
        "智能可穿戴设备",
        "智能机器人",
        "通用人工智能",
    ],
    "数字创意": [
        "动漫设计",
        "游戏设计",
        "电竞直播",
        "设计服务",
        "数字传媒",
        "会议会展",
    ],
    "氢能": [
        "氢能",
    ],
    "电磁能": [
        "电磁装备制造",
        "高速轨道交通",
    ],
    "量子科技": [
        "量子通信",
        "量子计算",
        "量子测量",
    ],
    "超级计算": ["超级计算"],
    "脑科学和类脑科学": [
        "类脑计算",
        "脑疾病诊疗",
        "脑机接口",
    ],
    "深地深海深空": [
        "深地探测",
        "深海探测",
        "深空观测",
    ],
}

from functools import lru_cache


@lru_cache(maxsize=256)
def get_large_cls_by_name(name):
    """
    通过子类获取大类名
    """
    for large_name, arr in wuhan_cls_struct.items():
        if name in arr:
            return large_name


# BERT 分类需要idx，将类别名转换为idx
def get_idx_by_name(name):
    for idx, (llm_cls, raw_cls) in enumerate(zip(cls_name, raw_attr)):
        if name == llm_cls or name == raw_cls:
            return idx


prompt_cls = f"""
你是一位企业分类的专家。请仔细阅读企业信息：industry_info  
接下来，请从下述类别列表中选出最符合该企业信息的单一类别，并返回该类别名作为列表。  
若企业信息不属于下述任何类别，请返回['其他']。  
类别列表: {cls_name}  
返回格式：['类别名']  
"""

industry_info_template = "{name}:\n经营范围:{business};\n大类名称:{large};\n中类名称:{mid};\n小类名称:{small};"


def trans_data(row):
    business = row.get("经营范围")
    name = row.get("企业名称")
    large = row.get("大类名称")
    mid = row.get("中类名称")
    small = row.get("小类名称")
    if not pd.isnull(name) and isinstance(name, str) and name.strip() != "":
        return industry_info_template.format(
            name=name, business=business, large=large, mid=mid, small=small
        )
    return ""


if __name__ == "__main__":

    print(prompt_cls)

    print(get_large_cls_by_name("深地探测"))

    print(sum(len(v) for v in wuhan_cls_struct.values()))

    # path = r"/home/jie/gitee/wuhan_industry_classify/dwuhan_cls_structata/new_data/武汉市产业门类和细分领域.xlsx"
    # df = pd.read_excel(path, header=2)
    # print(df.columns)
    # print(df.head())
    # print(df.shape)

    # print(df["细分领域"].values)

    # 观察是否 match
    # for idx in range(max(len(cls_name), len(raw_attr))):
    #     print(cls_name[idx], raw_attr[idx])
